import networkx as nx
import time
import json
from concurrent import futures
import multiprocessing

rsws = set()
routes = {}


def buildGraph():
    G = nx.Graph()
    i = 1
    for line in open('pnb_circuits'):
        arr = line.split()
        sw1_arr = arr[1].split(':')
        sw2_arr = arr[2].split(':')
        sw1 = sw1_arr[0]
        port1 = sw1_arr[1]
        sw2 = sw2_arr[0]
        port2 = sw2_arr[1]
        if sw1[0:3] == 'rsw':
            rsws.add(sw1)
        if sw2[0:3] == 'rsw':
            rsws.add(sw2)

        if not G.has_node(sw1):
            G.add_node(sw1, ip=(i << 16))
            i += 1
        if not G.has_node(sw2):
            G.add_node(sw2, ip=(i << 16))
            i += 1
        # o = G.number_of_edges()
        G.add_edge(sw1, sw2, link={sw1: port1, sw2: port2})
    return G


def worker(source):
    paths = nx.single_source_shortest_path(G, source)
    out = {}
    for j in rsws:
        path = paths[j]
        if len(path) > 1:
            out[G.nodes[j]['ip']] = G.get_edge_data(path[0], path[1])['link'][path[0]]
            routes[source][G.nodes[j]['ip']] = G.get_edge_data(path[0], path[1])['link'][path[0]]
    return out


if __name__ == '__main__':
    G = buildGraph()

    print(time.time())
    print(len(rsws))
    for sw in G.nodes:
        routes[sw] = {}
    start = time.time()

    with futures.ProcessPoolExecutor(max_workers=multiprocessing.cpu_count()) as executor:
        for node, route in zip(G.nodes, executor.map(worker, G.nodes)):
            routes[node] = route

        executor.shutdown(wait=True)
        with open('fbdata.txt', 'w') as f:
            json.dump(routes, f, indent=4)
        end = time.time()
        print(end - start)
